﻿

#include "utilityCommunicationManager.hpp"

namespace UTILITY {

typedef CommunicationManager *(*CommunicationManagerCreatePtr)(int *argc,
                                                               char ***argv);

Table<Word, CommunicationManagerCreatePtr> CommunicationManager::selectingTable;

namespace {

CommunicationManager *createCommunicationManager(int *argc, char ***argv) {
  return new CommunicationManager(argc, argv);
}

typedef pair<Word, CommunicationManagerCreatePtr> SelectEntry;

class AddCommunicationManagerToSelectingTable {
  const Word name_;

 public:
  AddCommunicationManagerToSelectingTable(const Word &name,
                                          CommunicationManagerCreatePtr ptr)
      : name_(name) {
    if (CommunicationManager::selectingTable.find(name_) ==
        CommunicationManager::selectingTable.end())
      CommunicationManager::selectingTable.insert(SelectEntry(name_, ptr));
  }

  ~AddCommunicationManagerToSelectingTable() {
    if (CommunicationManager::selectingTable.find(name_) !=
        CommunicationManager::selectingTable.end())
      CommunicationManager::selectingTable.erase(name_);
  }
} AddCommunicationManagerToSelectingTable_register("MPI",
                                                   createCommunicationManager);

}  // namespace

}  // namespace UTILITY

namespace UTILITY {

CommunicationManager::CommunicationManager(int *argc, char ***argv)
    : cDummyOut_(new std::ostream(0)) {
  int tag;
  MPI_Initialized(&tag);
  if (tag == 0) MPI_Init(argc, argv);

  MPI_Comm_rank(MPI_COMM_WORLD, &globalRank_);
  MPI_Comm_size(MPI_COMM_WORLD, &globalSize_);

  globalComm_ = new Communicator("globalComm");

  {
    MPI_Comm_dup(MPI_COMM_WORLD, &(globalComm_->comm_));
    globalComm_->rank_ = globalRank_;
    globalComm_->size_ = globalSize_;

    globalComm_->messageNum_ = 0;
  }
}

void CommunicationManager::createLogFile(bool TorF) {
  if (globalComm_->getMyId() == 0)
    globalComm_->log_ = new OStream(cout.rdbuf());
  else
    globalComm_->log_ = new DummyOStream();
  if (TorF) {
    if (globalComm_->getMyId() == 0 && access("logs", F_OK) != 0) {
#ifndef _WIN32
      mkdir("logs", S_IRWXU);
#else
      mkdir("logs");
#endif
    }
    if (globalComm_->getMyId() == 0) remakeDirectory("logs/comm.log");
    globalComm_->barrier();

    if (globalComm_->getMyId() == 0) remakeDirectory("logs/process.log");
    globalComm_->barrier();
    Word procLogFileName =
        "logs/process.log/" + to_string(globalComm_->getMyId());
    pOut_ = new OStream(procLogFileName);

    StrBuf *coutStrBuf = cout.rdbuf();
    pErr_ = new MultiOStream(&coutStrBuf);
    ((MultiOStream *)pErr_)->addBuffer(pOut_->getStrBuf());
  } else {
    pOut_ = new DummyOStream();
    StrBuf *coutStrBuf = cout.rdbuf();
    pErr_ = new MultiOStream(&coutStrBuf);
  }
}

CommunicationManager::~CommunicationManager() {
  delete globalComm_;
  delete pOut_;
  delete pErr_;
  delete cDummyOut_;
  int tag;
  MPI_Finalized(&tag);
  if (tag == 0) MPI_Finalize();
}

int CommunicationManager::duplicate(const Communicator &refComm,
                                    Communicator &newComm) {
  if (newComm.comm_ != MPI_COMM_NULL) {
    cerr << __FILE__ << " + " << __LINE__ << ": " << endl
         << __FUNCTION__ << ": " << endl
         << "Error: reconstruct a non-empty communicator!" << endl;
    exit(-1);
  }

  MPI_Comm_dup(refComm.comm_, &(newComm.comm_));

  newComm.size_ = refComm.size_;
  newComm.rank_ = refComm.rank_;
  return 0;
}

int CommunicationManager::split(const Communicator &refComm,
                                Communicator &newComm, const int color,
                                const unsigned int nPart, const int priority) {
  if (newComm.comm_ != MPI_COMM_NULL) {
    cerr << __FILE__ << " + " << __LINE__ << ": " << endl
         << __FUNCTION__ << ": " << endl
         << "Error: reconstruct a non-empty communicator!" << endl;
    exit(-1);
  }

  vector<int> *colors = new vector<int>(refComm.getMySize());
  (const_cast<Communicator &>(refComm))
      .allGather("allGatherColorsForSplit", &color, sizeof(int),
                 &((*colors)[0]), sizeof(int));
  (const_cast<Communicator &>(refComm)).finishTask("allGatherColorsForSplit");
  set<int> *colorSet = new set<int>;
  for (unsigned int i = 0; i < colors->size(); i++)
    colorSet->insert((*colors)[i]);
  if (colorSet->size() != nPart) {
    cerr << __FILE__ << " + " << __LINE__ << ": " << endl
         << __FUNCTION__ << ": " << endl
         << "Error: the colors for split is not consist with given part number!"
         << endl;
    exit(-1);
  }
  delete colors;
  delete colorSet;

  MPI_Comm_split(refComm.comm_, color, priority, &(newComm.comm_));
  setup(newComm);
  return 0;
}

int CommunicationManager::extract(const Communicator &refComm,
                                  Communicator &newComm, const bool color,
                                  int priority) {
  if (newComm.comm_ != MPI_COMM_NULL) {
    cerr << __FILE__ << " + " << __LINE__ << ": " << endl
         << __FUNCTION__ << ": " << endl
         << "Error: reconstruct a non-empty communicator!" << endl;
    exit(-1);
  }

  int color_int;
  if (color)
    color_int = 1;
  else
    color_int = 0;
  MPI_Comm_split(refComm.comm_, color_int, priority, &(newComm.comm_));

  if (!color)
    MPI_Comm_free(&(newComm.comm_));
  else
    setup(newComm);
  return 0;
}

int CommunicationManager::reBalance(const Communicator &oldComm,
                                    Communicator &newComm, const int weight,
                                    const float thresholdRatio) {
  cerr << __FILE__ << " +" << __LINE__ << ":" << endl
       << __FUNCTION__ << ":" << endl
       << "Error: reBalance is not implemented! " << endl;
  abort();
}

float CommunicationManager::calculateBalance(const Communicator &refComm,
                                             const int weight) {
  cerr << __FILE__ << " +" << __LINE__ << ":" << endl
       << __FUNCTION__ << ":" << endl
       << "Error: calculateBalance is not implemented! " << endl;
  abort();
}

}  // namespace UTILITY
